{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol1_features = np.random.rand(n_datapoints, n_features1)\n",
    "mol2_features = np.random.rand(n_datapoints, n_features2)\n",
    "x_ds = np.hstack([mol1_features, mol2_features])\n",
    "mol1s = [List of control molecules]\n",
    "mol2s = [List of target molecules]\n",
    "ys = np.random.rand(n_datapoints, n_tasks)\n",
    "datapoints1 = [MoleculeDatapoint(mol, y, x_d=x_d) for mol, y, x_d in zip(mol1s, ys, x_ds)]\n",
    "datapoints2 = [MoleculeDatapoint(mol) for mol in mol2s]\n",
    "dataset1 = MoleculeDataset(datapoints1)\n",
    "dataset2 = MoleculeDataset(datapoints2)\n",
    "multi_dataset = MulticomponentDataset(datasets=[dataset1, dataset2])\n",
    "dataloader = chemprop.data.build_dataloader(multi_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chemprop\n",
    "import rdkit\n",
    "import torch\n",
    "\n",
    "mol = rdkit.Chem.MolFromSmiles(\"CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C\")  # Penicillin\n",
    "dp = chemprop.data.MoleculeDatapoint(mol)\n",
    "ds = chemprop.data.MoleculeDataset([dp])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22.5 µs ± 1.33 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n",
      "66.8 µs ± 582 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n",
      "342 ns ± 10.6 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit list(mol.GetBonds())\n",
    "n_atoms = mol.GetNumAtoms()\n",
    "%timeit [mol.GetBondBetweenAtoms(u, v) for u in range(n_atoms) for v in range(u+1, n_atoms)]\n",
    "bond = mol.GetBondBetweenAtoms(0, 1)\n",
    "%timeit u, v = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "340 µs ± 32.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
      "455 ns ± 6.35 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1,  1,  2,  1,  5,  1, 22,  2,  3,  2, 19,  3,  4,  3,  7,\n",
       "         4,  5,  4,  6,  6,  7,  6,  9,  7,  8,  9, 10, 10, 11, 10, 12,\n",
       "        12, 13, 13, 14, 13, 18, 14, 15, 15, 16, 16, 17, 17, 18, 19, 20,\n",
       "        19, 21],\n",
       "       [ 1,  0,  2,  1,  5,  1, 22,  1,  3,  2, 19,  2,  4,  3,  7,  3,\n",
       "         5,  4,  6,  4,  7,  6,  9,  6,  8,  7, 10,  9, 11, 10, 12, 10,\n",
       "        13, 12, 14, 13, 18, 13, 15, 14, 16, 15, 17, 16, 18, 17, 20, 19,\n",
       "        21, 19]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Original\n",
    "%timeit ds[0]\n",
    "ds.cache = True\n",
    "%timeit ds[0]\n",
    "ds[0][0].edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "263 µs ± 15.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
      "439 ns ± 8.86 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1,  1,  2,  2,  3,  3,  4,  4,  5,  4,  6,  6,  7,  7,  8,\n",
       "         6,  9,  9, 10, 10, 11, 10, 12, 12, 13, 13, 14, 14, 15, 15, 16,\n",
       "        16, 17, 17, 18,  2, 19, 19, 20, 19, 21,  1, 22,  5,  1,  7,  3,\n",
       "        18, 13],\n",
       "       [ 1,  0,  2,  1,  3,  2,  4,  3,  5,  4,  6,  4,  7,  6,  8,  7,\n",
       "         9,  6, 10,  9, 11, 10, 12, 10, 13, 12, 14, 13, 15, 14, 16, 15,\n",
       "        17, 16, 18, 17, 19,  2, 20, 19, 21, 19, 22,  1,  1,  5,  3,  7,\n",
       "        13, 18]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# GetBonds\n",
    "%timeit ds[0]\n",
    "ds.cache = True\n",
    "%timeit ds[0]\n",
    "ds[0][0].edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "287 µs ± 6.53 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
      "439 ns ± 17.4 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0,  1,  1,  2,  2,  3,  3,  4,  4,  5,  4,  6,  6,  7,  7,  8,\n",
       "         6,  9,  9, 10, 10, 11, 10, 12, 12, 13, 13, 14, 14, 15, 15, 16,\n",
       "        16, 17, 17, 18,  2, 19, 19, 20, 19, 21,  1, 22,  5,  1,  7,  3,\n",
       "        18, 13],\n",
       "       [ 1,  0,  2,  1,  3,  2,  4,  3,  5,  4,  6,  4,  7,  6,  8,  7,\n",
       "         9,  6, 10,  9, 11, 10, 12, 10, 13, 12, 14, 13, 15, 14, 16, 15,\n",
       "        17, 16, 18, 17, 19,  2, 20, 19, 21, 19, 22,  1,  1,  5,  3,  7,\n",
       "        13, 18]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# List comprehension\n",
    "%timeit ds[0]\n",
    "ds.cache = True\n",
    "%timeit ds[0]\n",
    "ds[0][0].edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = rdkit.Chem.MolFromSmiles(\"\")\n",
    "dp = chemprop.data.MoleculeDatapoint(mol)\n",
    "ds = chemprop.data.MoleculeDataset([dp])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Datum(mg=MolGraph(V=array([], shape=(0, 72), dtype=float64), E=array([], shape=(0, 14), dtype=float64), edge_index=array([], shape=(2, 0), dtype=int64), rev_edge_index=array([], dtype=int64)), V_d=None, x_d=None, y=nan, weight=1.0, lt_mask=None, gt_mask=None)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp = chemprop.nn.BondMessagePassing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dl = chemprop.data.build_dataloader(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "mg, *_ = batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([], size=(0, 300), grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mp(mg, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = rdkit.Chem.MolFromSmiles(\"\")\n",
    "dp = chemprop.data.MoleculeDatapoint(mol)\n",
    "ds = chemprop.data.MoleculeDataset([dp])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chemprop\n",
    "import rdkit\n",
    "import torch\n",
    "\n",
    "mol = rdkit.Chem.MolFromSmiles(\"CC1(C(N2C(S1)C(C2=O)NC(=O)CC3=CC=CC=C3)C(=O)O)C\")  # Penicillin\n",
    "dp = chemprop.data.MoleculeDatapoint(mol)\n",
    "ds = chemprop.data.MoleculeDataset([dp])\n",
    "dl = chemprop.data.build_dataloader(ds, batch_size=1)\n",
    "batch = next(iter(dl))\n",
    "mg, *_ = batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/utilities/migration/utils.py:56: The loaded checkpoint was produced with Lightning v2.2.2, which is newer than your current Lightning version: v2.2.1\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MPNN(\n",
       "  (message_passing): BondMessagePassing(\n",
       "    (W_i): Linear(in_features=86, out_features=300, bias=False)\n",
       "    (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
       "    (W_o): Linear(in_features=372, out_features=300, bias=True)\n",
       "    (dropout): Dropout(p=0.0, inplace=False)\n",
       "    (tau): ReLU()\n",
       "    (V_d_transform): Identity()\n",
       "    (graph_transform): GraphTransform(\n",
       "      (V_transform): Identity()\n",
       "      (E_transform): Identity()\n",
       "    )\n",
       "  )\n",
       "  (agg): MeanAggregation()\n",
       "  (bn): BatchNorm1d(300, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (predictor): RegressionFFN(\n",
       "    (ffn): MLP(\n",
       "      (0): Sequential(\n",
       "        (0): Linear(in_features=300, out_features=300, bias=True)\n",
       "      )\n",
       "      (1): Sequential(\n",
       "        (0): ReLU()\n",
       "        (1): Dropout(p=0.0, inplace=False)\n",
       "        (2): Linear(in_features=300, out_features=1, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (criterion): MSELoss(task_weights=[[1.0]])\n",
       "    (output_transform): UnscaleTransform()\n",
       "  )\n",
       "  (X_d_transform): Identity()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint_path = \"tests/data/example_model_v2_regression_mol.ckpt\"\n",
    "mpnn = chemprop.models.MPNN.load_from_checkpoint(checkpoint_path)\n",
    "mpnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0146, 0.0000, 0.0711,  ..., 0.1050, 0.1969, 0.0000],\n",
       "        [0.0836, 0.0000, 0.0000,  ..., 0.2042, 0.3485, 0.2322],\n",
       "        [0.0000, 0.0000, 0.1735,  ..., 0.1548, 0.2860, 0.0000],\n",
       "        ...,\n",
       "        [0.0000, 0.0398, 0.2541,  ..., 0.0511, 0.0206, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0603,  ..., 0.0074, 0.1134, 0.0000],\n",
       "        [0.0146, 0.0000, 0.0711,  ..., 0.1050, 0.1969, 0.0000]],\n",
       "       grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpnn.message_passing(mg, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0146, 0.0000, 0.0711,  ..., 0.1050, 0.1969, 0.0000],\n",
       "        [0.0836, 0.0000, 0.0000,  ..., 0.2042, 0.3485, 0.2322],\n",
       "        [0.0000, 0.0000, 0.1735,  ..., 0.1548, 0.2860, 0.0000],\n",
       "        ...,\n",
       "        [0.0000, 0.0398, 0.2541,  ..., 0.0511, 0.0206, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0603,  ..., 0.0074, 0.1134, 0.0000],\n",
       "        [0.0146, 0.0000, 0.0711,  ..., 0.1050, 0.1969, 0.0000]],\n",
       "       grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpnn.message_passing(mg, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchmetrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.randn(10, 5)\n",
    "b = torch.randn(10, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = torchmetrics.MeanSquaredError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = torch.randn(3, 5, requires_grad=True)\n",
    "target = torch.randn(3, 5)\n",
    "output = torch.nn.functional.mse_loss(input, target, reduction=\"none\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.8324e-01, 2.7586e+00, 3.8411e+00, 2.0319e+00, 4.2811e+00],\n",
       "        [1.0199e+00, 1.0741e+00, 6.2624e-02, 2.0293e-04, 6.3696e-02],\n",
       "        [1.4223e+00, 3.7240e-03, 1.6341e+00, 3.9651e+00, 1.1566e-01]],\n",
       "       grad_fn=<MseLossBackward0>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.load(\n",
    "    \"/home/knathan/chemprop/chemprop_training/regression/2024-08-15T16-27-37/model_0/checkpoints/best-3-metric=0.05.ckpt\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "a[\"state_dict\"][\"metrics.0.task_weights\"] = torch.tensor([1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('message_passing.W_i.weight',\n",
       "              tensor([[ 0.0706,  0.0169, -0.0455,  ..., -0.0773,  0.0912,  0.0133],\n",
       "                      [-0.0699, -0.0489, -0.0614,  ...,  0.0055, -0.0094,  0.0202],\n",
       "                      [-0.0587, -0.0157,  0.0461,  ...,  0.0138, -0.0762,  0.1000],\n",
       "                      ...,\n",
       "                      [-0.0399,  0.0559,  0.0477,  ..., -0.0595,  0.0645, -0.0370],\n",
       "                      [ 0.0699, -0.0749,  0.1031,  ...,  0.0249,  0.1032, -0.0333],\n",
       "                      [-0.0739,  0.0413, -0.0260,  ..., -0.0189,  0.0846,  0.0117]])),\n",
       "             ('message_passing.W_h.weight',\n",
       "              tensor([[ 0.0551,  0.0161, -0.0060,  ..., -0.0489, -0.0355, -0.0217],\n",
       "                      [ 0.0240, -0.0067,  0.0056,  ..., -0.0368,  0.0477, -0.0200],\n",
       "                      [ 0.0155, -0.0106, -0.0087,  ..., -0.0260,  0.0372,  0.0429],\n",
       "                      ...,\n",
       "                      [ 0.0057, -0.0094,  0.0189,  ..., -0.0023,  0.0232, -0.0370],\n",
       "                      [ 0.0052, -0.0285, -0.0501,  ..., -0.0123, -0.0070,  0.0459],\n",
       "                      [ 0.0319,  0.0398, -0.0087,  ...,  0.0438, -0.0205,  0.0247]])),\n",
       "             ('message_passing.W_o.weight',\n",
       "              tensor([[ 0.0425,  0.0366,  0.0216,  ..., -0.0351, -0.0504,  0.0159],\n",
       "                      [-0.0145, -0.0094, -0.0099,  ..., -0.0051,  0.0467,  0.0066],\n",
       "                      [-0.0282, -0.0460, -0.0296,  ...,  0.0072, -0.0088,  0.0094],\n",
       "                      ...,\n",
       "                      [ 0.0279, -0.0015, -0.0078,  ...,  0.0193,  0.0256,  0.0485],\n",
       "                      [-0.0271, -0.0500, -0.0454,  ...,  0.0166, -0.0324,  0.0140],\n",
       "                      [-0.0023, -0.0222, -0.0443,  ...,  0.0161, -0.0423,  0.0468]])),\n",
       "             ('message_passing.W_o.bias',\n",
       "              tensor([-1.1304e-02, -3.8473e-02, -1.7424e-02, -8.8046e-04, -3.1658e-02,\n",
       "                      -4.7836e-02, -2.5356e-02,  1.1592e-02,  8.0597e-03,  6.9263e-03,\n",
       "                       4.5297e-02,  3.8366e-02,  1.0627e-02,  2.7768e-02, -3.8746e-02,\n",
       "                      -1.9652e-02,  3.3963e-02,  1.1734e-02, -1.6012e-02, -1.6624e-02,\n",
       "                       4.3988e-02, -1.0231e-02, -2.1802e-02,  1.6776e-02, -3.8954e-02,\n",
       "                      -2.6632e-02, -2.1116e-02,  3.5229e-02, -2.5180e-02,  1.3597e-02,\n",
       "                      -4.3012e-02, -5.6692e-05, -4.7552e-02, -4.8673e-02,  5.3184e-02,\n",
       "                      -2.2254e-02, -4.2718e-02,  4.6024e-02,  2.4164e-02,  2.6926e-02,\n",
       "                       4.3257e-03,  2.4133e-02, -3.2546e-02,  2.9479e-02, -3.4976e-02,\n",
       "                       6.6779e-03,  4.0294e-02,  1.5311e-02, -4.6442e-02, -5.2528e-03,\n",
       "                      -1.8169e-02,  5.4332e-02, -5.4104e-02, -1.3129e-02,  2.8221e-03,\n",
       "                       4.6811e-02, -1.9218e-02,  3.4289e-02,  4.9204e-03,  2.3035e-03,\n",
       "                       4.6189e-02,  3.2742e-03, -3.3104e-02,  2.1613e-02,  5.8889e-03,\n",
       "                       5.5967e-02, -5.3797e-02,  4.3298e-02,  3.5083e-02,  1.1211e-02,\n",
       "                       2.9126e-02,  2.3772e-02,  2.9378e-02,  2.2699e-02,  1.8136e-02,\n",
       "                      -4.1031e-02,  6.2324e-03,  3.1241e-02,  7.9547e-03,  2.3979e-02,\n",
       "                       2.6000e-02, -9.8677e-04, -6.9411e-03,  4.2141e-02,  1.4352e-02,\n",
       "                       3.7165e-02,  5.0779e-02, -2.5772e-02,  4.1488e-02,  2.8357e-02,\n",
       "                      -3.9428e-02, -4.0489e-02,  3.2342e-02, -4.6158e-02, -2.9769e-02,\n",
       "                      -7.2064e-03, -1.8967e-02,  2.2186e-02, -5.2161e-03,  6.0329e-03,\n",
       "                      -3.0137e-02, -5.1721e-03, -3.0214e-02,  4.7589e-02,  3.6432e-02,\n",
       "                       2.2697e-02, -1.0168e-02, -1.0288e-02, -4.0243e-02, -9.7010e-03,\n",
       "                       2.4228e-02, -4.6857e-02,  1.2850e-02,  4.3094e-02,  7.9320e-03,\n",
       "                      -2.1361e-02, -4.5189e-02,  1.8063e-02,  3.3066e-02,  2.3271e-02,\n",
       "                      -2.9678e-02,  1.6857e-02, -1.1120e-02, -1.3678e-02,  5.4269e-03,\n",
       "                       4.2425e-02,  2.1299e-02, -2.1539e-02,  3.0877e-02,  6.6988e-03,\n",
       "                       3.5534e-02, -9.5376e-03,  1.2235e-02, -1.1833e-02,  4.2716e-02,\n",
       "                      -7.0880e-03,  2.6442e-02,  4.6586e-02,  1.1214e-02,  2.1622e-02,\n",
       "                       7.1415e-03, -7.4622e-03, -4.7018e-02, -5.1831e-03, -3.9560e-02,\n",
       "                       4.4834e-02,  1.8867e-02, -3.5004e-02, -1.5880e-03,  1.5332e-02,\n",
       "                      -2.6821e-03, -3.8278e-02, -3.5086e-02,  3.3133e-03, -1.0369e-02,\n",
       "                      -1.4473e-02, -2.0081e-02, -1.8215e-02,  1.2375e-03,  3.9355e-02,\n",
       "                      -3.6434e-02, -1.0518e-02,  1.6295e-02, -1.0641e-02,  2.1844e-02,\n",
       "                       8.4319e-03,  6.8855e-03,  1.1405e-02, -1.5897e-03, -3.8670e-03,\n",
       "                       9.7160e-03, -1.6790e-02,  2.8037e-02, -4.7424e-02, -3.4596e-02,\n",
       "                      -1.1802e-02,  8.3450e-03, -3.1709e-02, -1.7009e-02, -3.7766e-02,\n",
       "                      -1.8032e-02, -2.9350e-02,  3.0843e-02,  7.4973e-03,  3.7044e-02,\n",
       "                      -2.5451e-02,  2.2911e-02, -6.5310e-03, -4.9357e-02, -2.3801e-02,\n",
       "                       2.9385e-02, -3.5699e-02,  3.1581e-02, -2.8849e-02, -1.6093e-02,\n",
       "                      -4.0973e-02, -2.4815e-02, -1.1365e-02,  2.1483e-02,  9.6250e-03,\n",
       "                       4.8921e-02, -2.9012e-02,  3.7005e-02,  4.4885e-02,  3.1444e-02,\n",
       "                      -4.0522e-02, -4.4704e-02,  3.5212e-02,  2.1225e-02,  5.3150e-02,\n",
       "                      -4.7975e-02,  3.3201e-02,  1.3442e-03, -3.5471e-02, -5.5576e-03,\n",
       "                       2.9601e-02,  4.8740e-03, -1.5098e-02, -1.9608e-02,  1.7828e-02,\n",
       "                       3.2324e-02,  1.8230e-03,  1.2947e-02, -6.0192e-03,  7.7552e-03,\n",
       "                      -6.6392e-03,  3.1892e-02, -1.4390e-02,  2.8080e-02, -1.7990e-02,\n",
       "                       2.9420e-02,  4.5269e-02, -2.0733e-02,  1.5052e-02,  3.0219e-02,\n",
       "                      -4.6810e-02, -2.8933e-02, -1.8122e-02, -3.3021e-02, -3.5576e-02,\n",
       "                       2.8806e-02,  3.4000e-02, -3.7130e-02,  2.7876e-02, -2.7954e-03,\n",
       "                       2.5227e-02, -5.3605e-02,  3.5885e-02, -3.5758e-02, -1.0221e-02,\n",
       "                      -2.9185e-02, -1.7799e-02, -1.9588e-02, -4.2681e-02, -2.8088e-02,\n",
       "                       3.2226e-02, -3.8566e-02,  1.7043e-02, -9.5702e-03,  1.5042e-02,\n",
       "                      -1.4779e-02,  3.3030e-02, -8.1090e-03, -6.6562e-03,  1.1123e-02,\n",
       "                      -4.5672e-03, -3.7714e-02, -4.7956e-02,  1.9636e-02, -4.9298e-02,\n",
       "                       3.2623e-02, -2.3274e-02, -2.4118e-02, -7.0314e-03, -1.4445e-02,\n",
       "                      -2.6205e-02, -4.8336e-02, -4.7223e-02, -8.8984e-03,  2.5293e-02,\n",
       "                      -4.1981e-02,  4.6695e-02, -5.1955e-03, -3.9184e-02, -2.2235e-03,\n",
       "                       5.0228e-02,  3.0144e-02,  1.6190e-02, -2.8378e-02, -1.4032e-02,\n",
       "                       8.7688e-03, -3.0148e-02, -2.4133e-03,  1.7978e-02, -3.1840e-02,\n",
       "                      -4.5051e-02,  5.1696e-02, -4.8709e-02, -4.7937e-02, -5.1205e-02])),\n",
       "             ('bn.weight',\n",
       "              tensor([1.0033, 0.9968, 1.0059, 1.0035, 1.0008, 1.0010, 1.0000, 0.9960, 0.9972,\n",
       "                      1.0008, 1.0026, 0.9972, 1.0000, 0.9954, 0.9952, 1.0048, 0.9978, 1.0016,\n",
       "                      0.9973, 1.0011, 0.9961, 0.9976, 0.9990, 0.9990, 0.9969, 0.9999, 1.0006,\n",
       "                      0.9974, 0.9976, 0.9963, 0.9970, 1.0037, 1.0000, 0.9984, 1.0045, 1.0012,\n",
       "                      0.9970, 1.0039, 1.0003, 0.9985, 1.0021, 0.9963, 0.9980, 0.9981, 1.0022,\n",
       "                      0.9989, 1.0005, 1.0002, 0.9987, 0.9984, 0.9969, 1.0023, 0.9966, 0.9963,\n",
       "                      0.9979, 1.0017, 1.0003, 1.0011, 0.9974, 0.9938, 0.9988, 0.9968, 1.0057,\n",
       "                      1.0023, 0.9967, 1.0063, 0.9966, 0.9962, 0.9990, 0.9948, 0.9976, 0.9987,\n",
       "                      0.9965, 1.0035, 1.0003, 0.9989, 1.0007, 1.0024, 0.9974, 0.9974, 1.0007,\n",
       "                      1.0006, 1.0000, 0.9955, 0.9952, 1.0007, 1.0030, 1.0043, 0.9977, 1.0009,\n",
       "                      1.0039, 1.0064, 0.9964, 0.9950, 1.0016, 0.9982, 1.0000, 1.0008, 1.0000,\n",
       "                      0.9975, 1.0021, 1.0006, 0.9983, 0.9981, 0.9958, 0.9956, 0.9962, 1.0007,\n",
       "                      0.9960, 1.0023, 1.0006, 0.9970, 1.0019, 1.0010, 1.0023, 1.0017, 0.9973,\n",
       "                      0.9956, 1.0048, 1.0018, 0.9996, 1.0007, 0.9993, 0.9947, 0.9976, 0.9963,\n",
       "                      1.0005, 0.9974, 1.0054, 1.0000, 0.9968, 0.9952, 1.0001, 1.0014, 1.0020,\n",
       "                      1.0009, 0.9965, 0.9969, 0.9995, 0.9970, 1.0023, 0.9982, 0.9956, 0.9977,\n",
       "                      0.9967, 1.0020, 0.9989, 0.9973, 0.9963, 1.0011, 0.9991, 1.0021, 1.0039,\n",
       "                      0.9972, 0.9959, 0.9960, 1.0016, 1.0021, 1.0049, 0.9995, 1.0022, 1.0017,\n",
       "                      1.0033, 0.9985, 0.9963, 0.9993, 0.9983, 0.9987, 0.9969, 0.9971, 0.9976,\n",
       "                      1.0025, 1.0017, 0.9981, 0.9987, 0.9968, 0.9967, 1.0054, 1.0036, 0.9945,\n",
       "                      1.0005, 1.0031, 0.9981, 0.9981, 1.0023, 0.9952, 0.9944, 0.9977, 0.9983,\n",
       "                      0.9970, 1.0010, 0.9979, 0.9978, 0.9955, 0.9998, 0.9972, 0.9984, 0.9967,\n",
       "                      0.9994, 0.9961, 1.0000, 1.0044, 0.9996, 1.0021, 0.9957, 1.0034, 0.9979,\n",
       "                      0.9977, 0.9994, 0.9972, 0.9995, 0.9976, 1.0005, 1.0009, 0.9970, 1.0003,\n",
       "                      0.9990, 0.9992, 0.9957, 0.9978, 0.9971, 0.9954, 0.9970, 1.0013, 0.9957,\n",
       "                      1.0047, 1.0023, 1.0044, 1.0063, 1.0008, 1.0029, 1.0024, 0.9958, 0.9961,\n",
       "                      0.9954, 0.9964, 1.0025, 0.9953, 1.0007, 0.9967, 0.9983, 0.9969, 0.9977,\n",
       "                      0.9959, 1.0030, 0.9975, 0.9965, 1.0034, 0.9983, 0.9963, 1.0023, 1.0001,\n",
       "                      0.9978, 1.0014, 0.9974, 1.0031, 1.0036, 1.0027, 1.0024, 1.0024, 1.0000,\n",
       "                      1.0045, 0.9994, 1.0035, 0.9972, 1.0016, 1.0017, 0.9985, 0.9963, 1.0009,\n",
       "                      0.9975, 1.0015, 0.9960, 0.9976, 0.9957, 1.0016, 0.9989, 0.9984, 1.0034,\n",
       "                      1.0000, 0.9965, 1.0012, 1.0047, 0.9971, 1.0033, 1.0027, 0.9958, 0.9977,\n",
       "                      0.9963, 0.9979, 0.9992, 1.0003, 0.9968, 1.0044, 0.9976, 1.0008, 1.0037,\n",
       "                      0.9974, 1.0035, 0.9958])),\n",
       "             ('bn.bias',\n",
       "              tensor([ 2.9983e-03,  3.2167e-03, -2.9835e-03, -2.9294e-03,  3.0069e-03,\n",
       "                       4.0542e-03,  3.1591e-03,  3.0537e-03,  3.5588e-03, -1.2295e-03,\n",
       "                      -3.0475e-03,  7.8301e-04,  5.0617e-03,  2.5663e-03,  4.2834e-03,\n",
       "                      -2.4629e-03,  1.5465e-03, -1.3940e-03,  2.7063e-03, -7.4029e-04,\n",
       "                       4.5405e-03, -3.5702e-03,  4.5513e-03, -1.9435e-03,  4.0263e-03,\n",
       "                       2.2946e-03, -1.4640e-03, -2.8162e-03,  2.7376e-03,  1.7759e-03,\n",
       "                      -4.0492e-04,  2.1472e-03, -7.7493e-04,  3.0055e-03, -2.4081e-03,\n",
       "                       1.3075e-03,  3.6234e-03, -2.6196e-03, -2.9037e-03,  3.6316e-03,\n",
       "                      -1.2183e-03,  4.5801e-04, -2.9326e-04, -2.8512e-03,  1.1984e-03,\n",
       "                       4.5528e-03, -2.9434e-03,  2.9423e-03, -5.8578e-04,  2.6783e-03,\n",
       "                       2.4301e-03, -2.1957e-03,  2.5244e-03, -1.7959e-03,  2.8836e-03,\n",
       "                      -3.0603e-03,  1.0545e-03, -2.0830e-03,  3.2607e-03, -5.3286e-04,\n",
       "                       2.1268e-03, -2.2501e-03, -3.0032e-03, -2.3702e-03,  4.2669e-03,\n",
       "                      -5.2014e-03,  3.5815e-03,  1.8626e-03, -3.6246e-03,  3.4647e-03,\n",
       "                       1.9979e-03,  1.1408e-03, -5.1387e-03,  3.6297e-03,  2.8005e-03,\n",
       "                       2.4614e-04,  4.0864e-03, -2.7910e-03, -1.7223e-03, -1.6560e-03,\n",
       "                      -2.3348e-03, -1.8964e-03,  5.4837e-03,  7.3555e-05, -3.9442e-03,\n",
       "                      -3.1029e-03,  3.5733e-03,  3.0850e-03, -5.5965e-03,  2.7398e-03,\n",
       "                      -4.3090e-03,  1.0782e-03,  1.5697e-03,  1.1498e-03,  3.8916e-04,\n",
       "                       3.5909e-03, -7.7310e-04, -2.1754e-03,  2.2949e-03,  2.2229e-03,\n",
       "                       2.5733e-04,  2.9483e-03, -2.8800e-03,  2.7101e-03, -9.7521e-04,\n",
       "                       4.0899e-03,  2.9202e-03,  1.3233e-03,  3.3375e-03, -3.3362e-03,\n",
       "                       1.4651e-03,  3.5404e-03, -1.3573e-03, -7.2574e-04, -2.9630e-03,\n",
       "                      -1.9181e-03,  3.0345e-06, -5.4027e-04, -7.0165e-04, -7.4072e-04,\n",
       "                       3.8194e-03,  5.7395e-03, -5.9579e-05, -4.5368e-03,  2.5730e-03,\n",
       "                      -1.0896e-03, -2.6944e-03, -8.5727e-04,  3.1955e-04, -8.1914e-05,\n",
       "                       3.8005e-03,  5.7537e-03, -3.8666e-03, -2.8832e-03,  4.3115e-03,\n",
       "                      -3.2916e-03,  1.1225e-03,  4.5203e-03, -5.6318e-05,  2.5502e-03,\n",
       "                      -3.8457e-03, -5.5145e-03,  4.9555e-03,  3.1363e-03,  2.7940e-03,\n",
       "                      -1.0560e-03, -9.5574e-04,  3.6676e-03,  2.2327e-03, -4.2050e-03,\n",
       "                      -1.7371e-03, -2.4932e-03, -2.0429e-03, -3.3904e-03,  3.3495e-03,\n",
       "                       4.9069e-03, -2.8416e-03, -2.1537e-03, -3.3344e-04,  7.6021e-04,\n",
       "                       2.2236e-03,  1.9696e-03, -1.9703e-03,  3.2538e-03,  1.2970e-03,\n",
       "                      -2.1817e-03,  4.5509e-03,  3.9202e-03,  3.7965e-03,  2.9237e-03,\n",
       "                       1.6992e-03, -2.5999e-03, -2.0205e-03,  2.9393e-03, -3.0546e-03,\n",
       "                       4.4215e-03,  3.4250e-03,  2.6873e-03,  3.4883e-03, -2.3163e-04,\n",
       "                      -2.4987e-03, -2.9119e-03,  2.4910e-03, -2.8796e-03, -3.7445e-03,\n",
       "                       3.5649e-03,  2.4095e-03,  4.2149e-03,  2.5730e-03,  3.6239e-03,\n",
       "                      -1.1785e-03,  5.9404e-03,  3.8755e-03,  4.2085e-03,  2.4760e-03,\n",
       "                       3.2471e-03, -1.5741e-03, -2.7342e-04, -2.7134e-03,  4.5522e-03,\n",
       "                      -5.9142e-04, -2.2577e-03, -2.7067e-03, -3.4728e-03,  1.5881e-03,\n",
       "                       3.0700e-03,  6.4912e-04, -5.4998e-04,  3.8781e-03,  2.5446e-03,\n",
       "                      -2.9546e-03, -4.5815e-03,  1.6828e-03,  3.2790e-04, -2.8965e-03,\n",
       "                       1.6468e-03,  9.3454e-04,  2.1927e-03,  2.0156e-03, -1.6306e-03,\n",
       "                       9.0565e-04,  3.0056e-03,  2.8142e-03, -3.7763e-03,  2.0348e-03,\n",
       "                      -1.7997e-03,  6.1756e-04, -5.2506e-03, -4.8400e-03, -4.0161e-03,\n",
       "                       1.7103e-03, -2.9479e-03,  2.8712e-03,  3.1151e-03,  3.2711e-03,\n",
       "                      -1.0826e-04,  4.1681e-03, -3.5747e-03, -1.2852e-03,  3.7105e-03,\n",
       "                       2.6544e-03, -1.8356e-03,  2.4458e-03,  3.8169e-03, -1.7801e-03,\n",
       "                      -1.8200e-03,  2.6918e-03,  5.7959e-04, -9.1062e-04,  4.0207e-03,\n",
       "                      -3.4336e-03, -2.3241e-03,  2.2996e-03, -1.9961e-03, -3.4619e-03,\n",
       "                      -1.9873e-03, -1.8049e-03, -2.4167e-03, -1.5379e-03,  7.2068e-04,\n",
       "                      -4.0260e-05,  5.3464e-03,  3.0172e-03,  3.6957e-03,  2.5006e-03,\n",
       "                       9.7741e-04, -2.6105e-03, -1.6638e-03,  2.9245e-03,  1.7587e-03,\n",
       "                      -2.0594e-03,  1.2118e-03,  3.0792e-03,  2.0757e-03,  4.5420e-03,\n",
       "                       3.9449e-04,  7.4594e-05, -3.3196e-03,  1.9300e-03,  1.2875e-03,\n",
       "                      -1.6727e-03, -1.5769e-05,  3.8349e-03,  3.8636e-04,  1.5103e-03,\n",
       "                       1.4692e-03,  3.0448e-03,  2.3456e-03,  1.3412e-03,  2.5942e-03,\n",
       "                      -1.5803e-04, -3.3179e-03,  3.7802e-03, -3.2245e-03,  4.3890e-04,\n",
       "                       1.2354e-04, -2.3398e-03,  3.9685e-03, -1.1293e-03,  4.1611e-03])),\n",
       "             ('bn.running_mean',\n",
       "              tensor([5.9996e-03, 6.0317e-03, 1.7883e-03, 1.1019e-02, 2.5023e-06, 8.0262e-03,\n",
       "                      0.0000e+00, 7.1190e-03, 6.4138e-04, 4.5117e-03, 2.1331e-02, 3.3860e-04,\n",
       "                      0.0000e+00, 5.7292e-02, 5.1030e-04, 5.5463e-04, 2.8915e-02, 2.3194e-02,\n",
       "                      3.4753e-02, 5.0719e-04, 3.0886e-02, 1.1589e-05, 1.3766e-03, 4.1617e-03,\n",
       "                      2.3505e-02, 3.0266e-03, 3.5453e-03, 9.6264e-03, 2.5043e-02, 1.0884e-02,\n",
       "                      5.7304e-03, 3.6388e-03, 0.0000e+00, 2.2798e-05, 2.2195e-03, 8.4241e-03,\n",
       "                      2.7551e-02, 1.4417e-03, 7.9673e-03, 6.5603e-03, 1.3227e-02, 1.4504e-02,\n",
       "                      2.9334e-06, 7.1140e-03, 4.5255e-03, 5.0127e-03, 8.3108e-03, 1.0050e-02,\n",
       "                      7.8695e-03, 3.6205e-05, 2.4952e-04, 3.6416e-02, 3.1782e-03, 1.1999e-04,\n",
       "                      1.8478e-03, 2.2093e-02, 2.1094e-03, 2.5969e-02, 3.5535e-04, 1.1791e-02,\n",
       "                      2.7098e-02, 1.5125e-06, 7.8110e-04, 3.1907e-02, 6.6545e-05, 6.3280e-03,\n",
       "                      1.3123e-03, 2.2047e-02, 2.5472e-02, 4.6569e-03, 4.0808e-06, 3.3070e-02,\n",
       "                      3.8427e-02, 1.8204e-03, 2.0470e-05, 6.5955e-03, 2.0194e-03, 5.9259e-03,\n",
       "                      3.2310e-02, 8.4303e-05, 2.3236e-03, 8.4918e-07, 7.2321e-03, 3.3940e-03,\n",
       "                      8.7707e-05, 2.2778e-02, 2.5069e-03, 6.6309e-04, 3.1649e-08, 2.4449e-02,\n",
       "                      4.7569e-03, 1.2646e-03, 3.8340e-03, 3.8062e-03, 1.0368e-02, 1.5561e-03,\n",
       "                      0.0000e+00, 1.9439e-02, 0.0000e+00, 6.8192e-04, 1.0841e-03, 1.3471e-02,\n",
       "                      4.8461e-04, 1.6337e-02, 3.9447e-02, 2.7003e-03, 2.9765e-05, 2.5389e-02,\n",
       "                      6.5246e-05, 1.5681e-02, 7.1821e-03, 1.3128e-02, 5.4884e-06, 9.2366e-03,\n",
       "                      1.4159e-02, 1.6008e-02, 1.6612e-05, 2.4477e-04, 4.8010e-03, 3.2323e-02,\n",
       "                      2.8059e-05, 3.3883e-03, 6.9262e-03, 2.3474e-04, 7.9364e-04, 1.4285e-02,\n",
       "                      5.6118e-02, 1.3957e-03, 3.9597e-03, 6.7442e-04, 4.9940e-04, 4.9245e-03,\n",
       "                      3.9650e-03, 7.2841e-03, 7.8904e-03, 4.1480e-02, 8.5537e-03, 2.8654e-02,\n",
       "                      1.7650e-05, 9.4302e-08, 2.5946e-02, 2.8268e-02, 3.0528e-03, 2.4470e-04,\n",
       "                      5.0952e-07, 2.4923e-03, 3.9301e-02, 6.2683e-05, 1.2753e-04, 1.7916e-02,\n",
       "                      8.4331e-03, 2.1609e-02, 7.1982e-03, 2.4126e-02, 7.3836e-04, 1.5527e-04,\n",
       "                      6.9376e-05, 2.9736e-04, 2.2342e-03, 3.6643e-07, 7.9125e-03, 1.0067e-02,\n",
       "                      4.8342e-03, 1.0556e-03, 1.0693e-05, 3.9808e-03, 7.8145e-03, 2.4616e-07,\n",
       "                      1.5339e-04, 3.0379e-04, 2.8216e-02, 1.6389e-03, 2.1170e-03, 2.9291e-05,\n",
       "                      1.2275e-02, 1.9028e-03, 1.8553e-03, 3.3427e-04, 1.3428e-02, 8.9350e-03,\n",
       "                      1.0540e-02, 6.5879e-05, 6.7317e-04, 2.3394e-02, 1.9554e-02, 2.7866e-03,\n",
       "                      1.4856e-02, 2.2277e-02, 2.1048e-03, 1.8054e-02, 3.2366e-02, 1.7324e-06,\n",
       "                      9.0987e-03, 7.7648e-03, 2.1840e-03, 2.3373e-04, 2.4214e-03, 1.5512e-02,\n",
       "                      1.8707e-02, 3.3140e-02, 2.5308e-02, 3.1334e-04, 5.4786e-03, 3.3509e-03,\n",
       "                      1.4868e-04, 3.1236e-03, 2.8015e-05, 5.2305e-03, 2.0788e-05, 1.4806e-02,\n",
       "                      7.1128e-07, 2.5907e-02, 1.3713e-02, 2.8416e-06, 2.7722e-04, 5.0476e-03,\n",
       "                      1.0848e-03, 1.1363e-03, 4.8644e-03, 2.8875e-09, 2.6244e-07, 1.1356e-02,\n",
       "                      1.5239e-03, 1.3667e-02, 2.9040e-02, 1.8037e-03, 2.4959e-02, 5.0413e-03,\n",
       "                      1.8661e-03, 1.9202e-02, 1.0083e-02, 4.4256e-02, 1.1194e-02, 4.7456e-03,\n",
       "                      2.6440e-02, 4.1307e-03, 1.2197e-03, 1.3870e-03, 1.7738e-02, 1.0732e-03,\n",
       "                      2.4542e-05, 1.0208e-02, 7.3856e-08, 5.5114e-03, 5.9488e-03, 1.3749e-06,\n",
       "                      9.7147e-05, 2.4719e-02, 2.3731e-06, 2.3523e-04, 1.4885e-02, 5.0340e-03,\n",
       "                      1.4846e-07, 8.1106e-03, 1.5237e-04, 7.0195e-03, 9.2342e-03, 3.0843e-03,\n",
       "                      6.5658e-03, 2.0747e-02, 0.0000e+00, 5.0696e-03, 2.2728e-02, 1.8662e-03,\n",
       "                      3.9860e-03, 5.7752e-03, 5.3912e-02, 1.2127e-04, 6.5313e-04, 2.6554e-03,\n",
       "                      4.3055e-07, 6.3130e-03, 2.6714e-07, 4.4304e-04, 2.0553e-03, 6.2089e-03,\n",
       "                      9.2664e-03, 1.1430e-02, 1.5972e-02, 1.6042e-03, 6.8243e-06, 8.4010e-03,\n",
       "                      6.0464e-03, 2.1622e-03, 1.2293e-02, 2.5390e-04, 3.8274e-03, 1.1309e-04,\n",
       "                      1.0329e-03, 1.7351e-04, 3.4152e-03, 2.2246e-02, 1.5364e-04, 1.2961e-02,\n",
       "                      2.4346e-02, 1.9170e-02, 2.9175e-02, 3.7505e-06, 5.5555e-03, 3.8724e-03])),\n",
       "             ('bn.running_var',\n",
       "              tensor([0.1853, 0.1853, 0.1853, 0.1854, 0.1853, 0.1854, 0.1853, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1855, 0.1853, 0.1853, 0.1869, 0.1853, 0.1853, 0.1857, 0.1856,\n",
       "                      0.1860, 0.1853, 0.1857, 0.1853, 0.1853, 0.1853, 0.1857, 0.1853, 0.1853,\n",
       "                      0.1854, 0.1859, 0.1854, 0.1854, 0.1853, 0.1853, 0.1853, 0.1853, 0.1855,\n",
       "                      0.1859, 0.1853, 0.1854, 0.1853, 0.1855, 0.1855, 0.1853, 0.1854, 0.1853,\n",
       "                      0.1853, 0.1854, 0.1854, 0.1854, 0.1853, 0.1853, 0.1858, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1855, 0.1853, 0.1858, 0.1853, 0.1855, 0.1857, 0.1853, 0.1853,\n",
       "                      0.1857, 0.1853, 0.1854, 0.1853, 0.1855, 0.1858, 0.1853, 0.1853, 0.1861,\n",
       "                      0.1859, 0.1853, 0.1853, 0.1854, 0.1853, 0.1854, 0.1858, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1854, 0.1853, 0.1853, 0.1857, 0.1853, 0.1853, 0.1853, 0.1855,\n",
       "                      0.1854, 0.1853, 0.1853, 0.1853, 0.1854, 0.1853, 0.1853, 0.1855, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1854, 0.1853, 0.1854, 0.1860, 0.1853, 0.1853, 0.1857,\n",
       "                      0.1853, 0.1855, 0.1853, 0.1855, 0.1853, 0.1854, 0.1857, 0.1855, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1858, 0.1853, 0.1853, 0.1854, 0.1853, 0.1853, 0.1855,\n",
       "                      0.1868, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853,\n",
       "                      0.1867, 0.1854, 0.1858, 0.1853, 0.1853, 0.1857, 0.1859, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1860, 0.1853, 0.1853, 0.1855, 0.1854, 0.1856, 0.1854,\n",
       "                      0.1856, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1854,\n",
       "                      0.1854, 0.1853, 0.1853, 0.1853, 0.1854, 0.1853, 0.1853, 0.1853, 0.1857,\n",
       "                      0.1853, 0.1853, 0.1853, 0.1855, 0.1853, 0.1853, 0.1853, 0.1855, 0.1854,\n",
       "                      0.1854, 0.1853, 0.1853, 0.1856, 0.1856, 0.1853, 0.1856, 0.1856, 0.1853,\n",
       "                      0.1856, 0.1858, 0.1853, 0.1853, 0.1854, 0.1853, 0.1853, 0.1853, 0.1855,\n",
       "                      0.1856, 0.1858, 0.1857, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1855, 0.1853, 0.1858, 0.1855, 0.1853, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1854, 0.1853, 0.1853, 0.1855, 0.1853, 0.1854, 0.1857,\n",
       "                      0.1853, 0.1856, 0.1854, 0.1853, 0.1855, 0.1854, 0.1862, 0.1854, 0.1854,\n",
       "                      0.1856, 0.1853, 0.1853, 0.1853, 0.1856, 0.1853, 0.1853, 0.1854, 0.1853,\n",
       "                      0.1854, 0.1853, 0.1853, 0.1853, 0.1856, 0.1853, 0.1853, 0.1854, 0.1853,\n",
       "                      0.1853, 0.1854, 0.1853, 0.1853, 0.1854, 0.1853, 0.1854, 0.1855, 0.1853,\n",
       "                      0.1853, 0.1855, 0.1853, 0.1853, 0.1853, 0.1868, 0.1853, 0.1853, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1853, 0.1853, 0.1853, 0.1854, 0.1854, 0.1854, 0.1855,\n",
       "                      0.1853, 0.1853, 0.1854, 0.1853, 0.1853, 0.1853, 0.1853, 0.1854, 0.1853,\n",
       "                      0.1853, 0.1853, 0.1853, 0.1855, 0.1853, 0.1854, 0.1857, 0.1855, 0.1857,\n",
       "                      0.1853, 0.1853, 0.1853])),\n",
       "             ('bn.num_batches_tracked', tensor(16)),\n",
       "             ('predictor.ffn.0.0.weight',\n",
       "              tensor([[ 0.0403, -0.0172,  0.0397,  ...,  0.0406,  0.0030,  0.0560],\n",
       "                      [ 0.0273,  0.0248,  0.0273,  ..., -0.0396,  0.0258,  0.0365],\n",
       "                      [-0.0419,  0.0047,  0.0206,  ...,  0.0538,  0.0179,  0.0272],\n",
       "                      ...,\n",
       "                      [-0.0347, -0.0483,  0.0033,  ...,  0.0483,  0.0183,  0.0314],\n",
       "                      [-0.0292, -0.0079,  0.0008,  ..., -0.0481, -0.0160,  0.0372],\n",
       "                      [-0.0241,  0.0386,  0.0014,  ...,  0.0357,  0.0128, -0.0053]])),\n",
       "             ('predictor.ffn.0.0.bias',\n",
       "              tensor([ 0.0432, -0.0104, -0.0075,  0.0185, -0.0549, -0.0337, -0.0026,  0.0473,\n",
       "                       0.0153,  0.0008,  0.0373,  0.0151, -0.0466, -0.0358,  0.0074,  0.0324,\n",
       "                      -0.0307,  0.0494,  0.0139, -0.0195, -0.0199,  0.0247, -0.0038,  0.0250,\n",
       "                       0.0055,  0.0476,  0.0162,  0.0582, -0.0214, -0.0395, -0.0031,  0.0247,\n",
       "                       0.0097, -0.0040,  0.0106,  0.0435,  0.0286,  0.0107, -0.0282, -0.0342,\n",
       "                      -0.0304,  0.0359, -0.0165,  0.0303, -0.0577, -0.0520, -0.0125, -0.0023,\n",
       "                      -0.0183,  0.0132, -0.0453, -0.0284,  0.0438, -0.0034,  0.0090, -0.0395,\n",
       "                      -0.0107, -0.0529, -0.0580, -0.0249, -0.0319,  0.0493,  0.0319,  0.0226,\n",
       "                       0.0463,  0.0163, -0.0143,  0.0207,  0.0301, -0.0449,  0.0499,  0.0068,\n",
       "                      -0.0589, -0.0429,  0.0345,  0.0229,  0.0183,  0.0026, -0.0523, -0.0207,\n",
       "                       0.0531,  0.0364, -0.0163,  0.0121, -0.0347, -0.0489, -0.0388, -0.0117,\n",
       "                      -0.0063,  0.0103, -0.0107, -0.0511,  0.0409,  0.0179, -0.0410, -0.0055,\n",
       "                      -0.0537, -0.0344, -0.0298, -0.0098, -0.0116, -0.0547, -0.0175,  0.0419,\n",
       "                      -0.0511, -0.0311,  0.0418,  0.0182, -0.0155,  0.0462,  0.0480,  0.0056,\n",
       "                      -0.0031,  0.0141,  0.0417, -0.0241,  0.0199,  0.0589, -0.0043,  0.0101,\n",
       "                       0.0400, -0.0225,  0.0458, -0.0221, -0.0077,  0.0153,  0.0300,  0.0496,\n",
       "                       0.0510, -0.0211, -0.0312,  0.0095,  0.0206, -0.0231, -0.0455,  0.0193,\n",
       "                       0.0041, -0.0296,  0.0159, -0.0320,  0.0103,  0.0406, -0.0302, -0.0249,\n",
       "                       0.0226, -0.0013, -0.0128,  0.0483, -0.0256, -0.0412, -0.0266,  0.0504,\n",
       "                      -0.0372, -0.0551,  0.0169, -0.0362,  0.0024,  0.0222,  0.0466,  0.0117,\n",
       "                      -0.0230,  0.0520,  0.0356,  0.0504, -0.0486, -0.0383, -0.0075, -0.0365,\n",
       "                      -0.0065, -0.0401, -0.0109, -0.0472,  0.0242,  0.0061,  0.0304, -0.0078,\n",
       "                      -0.0600,  0.0055, -0.0266, -0.0118, -0.0235,  0.0213,  0.0151,  0.0195,\n",
       "                      -0.0091, -0.0326, -0.0406,  0.0124, -0.0193, -0.0085, -0.0567, -0.0168,\n",
       "                      -0.0392,  0.0272, -0.0542, -0.0499,  0.0190, -0.0097,  0.0254,  0.0184,\n",
       "                      -0.0295, -0.0320,  0.0392,  0.0474, -0.0245,  0.0501, -0.0074, -0.0415,\n",
       "                       0.0534, -0.0277,  0.0549,  0.0259, -0.0600,  0.0001, -0.0009,  0.0031,\n",
       "                      -0.0013, -0.0152, -0.0370, -0.0007, -0.0384, -0.0510,  0.0124,  0.0506,\n",
       "                      -0.0371,  0.0436, -0.0195,  0.0177,  0.0230, -0.0347, -0.0394,  0.0477,\n",
       "                       0.0565, -0.0194, -0.0532, -0.0430,  0.0052, -0.0569, -0.0155,  0.0212,\n",
       "                      -0.0041, -0.0192,  0.0014,  0.0257, -0.0271, -0.0366,  0.0067, -0.0046,\n",
       "                       0.0409,  0.0298, -0.0469,  0.0273,  0.0442, -0.0089, -0.0063, -0.0177,\n",
       "                      -0.0184, -0.0093, -0.0377, -0.0175, -0.0082, -0.0320, -0.0130, -0.0488,\n",
       "                       0.0170, -0.0516,  0.0423,  0.0448,  0.0444, -0.0147, -0.0013, -0.0309,\n",
       "                      -0.0489,  0.0505,  0.0238, -0.0426,  0.0366,  0.0566,  0.0158, -0.0358,\n",
       "                       0.0176, -0.0197,  0.0438,  0.0393,  0.0238,  0.0276, -0.0062,  0.0236,\n",
       "                      -0.0291, -0.0269,  0.0348, -0.0488, -0.0522,  0.0219,  0.0186, -0.0603,\n",
       "                       0.0552,  0.0087,  0.0549,  0.0258])),\n",
       "             ('predictor.ffn.1.2.weight',\n",
       "              tensor([[-3.7370e-02,  1.2524e-02, -5.6026e-03,  5.1304e-02, -1.7573e-02,\n",
       "                       -1.1862e-02,  5.3584e-02, -4.3021e-02, -6.4778e-03, -2.5587e-02,\n",
       "                       -2.6129e-02, -6.3961e-03,  4.4003e-02,  4.0449e-02, -4.9500e-02,\n",
       "                       -3.7521e-02,  2.7940e-02,  4.7645e-02,  3.2447e-02, -3.9302e-03,\n",
       "                       -2.1370e-02,  1.6764e-02, -4.0053e-02, -3.1886e-02,  3.3077e-02,\n",
       "                        5.4420e-05,  2.7563e-02, -1.3385e-02, -3.3598e-02,  1.7829e-02,\n",
       "                       -5.0934e-03, -9.9870e-03,  2.8752e-02,  3.2577e-03, -3.6542e-03,\n",
       "                       -9.4377e-03,  4.3098e-02,  1.0824e-03, -1.9423e-02, -4.8490e-02,\n",
       "                        5.6046e-02, -4.3513e-02,  1.5941e-02, -2.5362e-02, -4.5007e-02,\n",
       "                        5.5456e-02, -3.5876e-02,  3.5563e-02,  3.3536e-02, -1.8639e-02,\n",
       "                       -4.0825e-02, -4.7310e-02, -4.2375e-02,  3.8371e-03, -5.4863e-02,\n",
       "                       -3.4437e-02, -4.1843e-02, -9.1780e-03, -7.4100e-03,  3.7572e-02,\n",
       "                        2.2665e-02, -5.2648e-02,  1.8550e-02,  5.9515e-03,  9.3363e-03,\n",
       "                       -3.6633e-02, -4.6670e-02,  9.4602e-03, -1.4803e-02, -3.0332e-02,\n",
       "                        5.1598e-02, -3.1526e-02,  4.2351e-02, -6.6361e-03,  4.8949e-03,\n",
       "                       -1.2248e-03,  1.3363e-02, -5.0477e-02,  2.4753e-02, -2.9339e-03,\n",
       "                       -2.6261e-02,  1.0329e-02,  2.6579e-03, -1.7709e-02,  1.7040e-02,\n",
       "                       -4.5151e-02,  1.2901e-02, -3.1480e-02,  4.5777e-02, -1.5395e-02,\n",
       "                       -4.2918e-02,  2.3874e-02,  4.2527e-02, -4.3337e-02,  3.5736e-02,\n",
       "                        4.7358e-02,  3.4818e-03,  2.7667e-02,  6.2228e-03, -1.1899e-03,\n",
       "                       -4.9705e-02, -3.4672e-02, -2.5415e-02,  2.2858e-03, -2.4991e-02,\n",
       "                       -1.6740e-02,  4.5004e-02, -1.1793e-02, -2.2895e-02, -1.2562e-02,\n",
       "                        2.6278e-02,  3.7152e-02,  5.2831e-02,  7.0276e-03, -3.2727e-02,\n",
       "                       -3.6156e-02,  5.4360e-02,  2.2302e-02,  4.8093e-02, -7.1987e-03,\n",
       "                       -3.6580e-02,  8.4851e-03, -3.7694e-02,  1.4725e-02, -5.5010e-02,\n",
       "                       -5.0632e-03,  4.0521e-02, -4.3414e-02, -2.8069e-02, -1.9888e-02,\n",
       "                        4.3080e-02, -3.0105e-02,  3.4326e-02, -3.7370e-02, -3.2847e-02,\n",
       "                       -1.1306e-02,  1.8656e-02, -4.9502e-04,  1.3088e-03, -2.2937e-02,\n",
       "                       -4.7985e-02, -3.7962e-02,  5.8178e-02,  5.2559e-02,  1.9985e-02,\n",
       "                        1.5631e-02, -1.6851e-02,  3.7614e-02, -4.8487e-02,  1.6453e-02,\n",
       "                        5.5260e-02,  4.2318e-02,  1.9642e-02, -3.0932e-02,  4.2697e-02,\n",
       "                       -2.3790e-02,  4.6843e-02, -2.8206e-02, -3.9776e-02,  6.1090e-03,\n",
       "                       -4.5110e-02, -1.4732e-02, -2.6706e-02,  5.4184e-02,  2.9133e-02,\n",
       "                        4.2807e-03,  8.4943e-03,  4.1638e-02,  5.5031e-02,  5.5712e-02,\n",
       "                        3.0725e-02, -4.4849e-03,  1.3911e-02, -1.8394e-03,  3.5050e-02,\n",
       "                       -3.0132e-02, -1.6487e-02,  1.6839e-02,  5.3991e-02,  6.3747e-03,\n",
       "                       -1.2767e-02,  1.4731e-02, -1.5415e-02, -2.6316e-02,  1.8080e-02,\n",
       "                        3.4356e-02, -3.0796e-02, -2.4189e-02,  2.5846e-03,  3.7984e-02,\n",
       "                       -4.6796e-02, -5.6010e-02,  3.9749e-02, -7.0712e-03, -7.4819e-03,\n",
       "                       -1.9106e-02,  5.2357e-02, -8.4244e-03,  2.8796e-02,  3.1351e-02,\n",
       "                       -5.1485e-02, -4.8105e-02,  2.7331e-02, -3.8832e-02,  1.2704e-03,\n",
       "                       -4.7566e-02,  4.8285e-02,  1.1345e-02,  2.2287e-04, -1.2127e-02,\n",
       "                        2.8947e-03, -5.3746e-02, -1.4423e-02, -4.7554e-02, -1.0545e-02,\n",
       "                        4.3091e-02,  1.8723e-02, -8.5750e-03, -4.8665e-02, -2.8907e-02,\n",
       "                        1.1600e-02,  3.2440e-02, -1.9376e-02,  1.0606e-02, -2.9975e-02,\n",
       "                       -1.7441e-02, -4.5052e-02, -5.3992e-02,  5.2634e-02, -5.2933e-02,\n",
       "                       -1.5525e-02, -6.4573e-03,  1.6406e-02,  5.1426e-02, -3.3195e-02,\n",
       "                        5.2496e-02,  4.0666e-02,  7.2251e-03, -1.5891e-02,  2.9241e-02,\n",
       "                        2.3724e-02,  3.0773e-02, -4.9185e-03, -2.9290e-02, -2.5603e-02,\n",
       "                        3.9914e-02,  5.2935e-02, -3.5327e-02,  2.6485e-02,  1.9386e-02,\n",
       "                       -3.3390e-02, -2.8076e-02,  3.6633e-02, -2.6471e-02, -5.5247e-02,\n",
       "                       -2.2887e-02, -8.8440e-03, -1.6074e-02,  1.5775e-02,  5.5354e-02,\n",
       "                        5.1264e-04, -2.2570e-02,  1.4420e-02,  8.1544e-03,  3.0691e-02,\n",
       "                       -5.2727e-02,  1.4428e-02,  4.7017e-02, -5.5784e-02,  3.4190e-02,\n",
       "                       -2.9872e-02,  2.6775e-02,  2.3817e-03,  6.6333e-04, -4.4412e-02,\n",
       "                        2.6779e-02,  1.5083e-02,  4.5320e-02, -3.5502e-02, -4.1737e-02,\n",
       "                       -3.4796e-02, -1.8898e-02, -4.1161e-02, -3.7673e-02, -1.3815e-02,\n",
       "                       -4.3093e-02,  4.2474e-02,  1.1878e-02, -2.9480e-03,  2.3156e-02,\n",
       "                       -3.7098e-02,  2.1611e-02, -2.1721e-03, -2.4970e-02, -2.8555e-02,\n",
       "                       -3.0127e-02,  3.8005e-02,  5.4677e-02, -5.0412e-02, -3.7432e-02]])),\n",
       "             ('predictor.ffn.1.2.bias', tensor([0.0336])),\n",
       "             ('predictor.criterion.task_weights', tensor([[1.]])),\n",
       "             ('predictor.output_transform.mean', tensor([[-3.0657]])),\n",
       "             ('predictor.output_transform.scale', tensor([[2.0944]])),\n",
       "             ('metrics.0.task_weights', tensor([[1.]])),\n",
       "             ('metrics.1.task_weights', tensor([[1.]])),\n",
       "             ('metrics.2.task_weights', tensor([[1.]])),\n",
       "             ('metrics.3.task_weights', tensor([[1.]])),\n",
       "             ('metrics.4.task_weights', tensor([[1.]]))])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[\"state_dict\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"activation\":       RELU\n",
       "\"cls\":              <class 'chemprop.nn.predictors.RegressionFFN'>\n",
       "\"criterion\":        None\n",
       "\"dropout\":          0.0\n",
       "\"hidden_dim\":       300\n",
       "\"input_dim\":        300\n",
       "\"n_layers\":         1\n",
       "\"n_tasks\":          1\n",
       "\"output_transform\": UnscaleTransform()\n",
       "\"task_weights\":     None\n",
       "\"threshold\":        None"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[\"hyper_parameters\"][\"predictor\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chemprop",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
